package main

import (
	"encoding/binary"
	"io"
	"log"
	"os"
	"syscall"

	"github.com/ceph/go-ceph/rados"
	"github.com/ceph/go-ceph/rbd"
	"golang.org/x/sys/unix"
)

const (
	nbdDevice = "/dev/nbd5"
	poolName  = "armcloud-dev"
	imageName = "ACN250710DL94RXJ"

	NBD_REQUEST_MAGIC = 0x25609513
	NBD_REPLY_MAGIC   = 0x67446698

	NBD_CMD_READ  = 0
	NBD_CMD_WRITE = 1
	NBD_CMD_DISC  = 2
	NBD_CMD_FLUSH = 3
	NBD_CMD_TRIM  = 4

	NBD_SET_SOCK        = 0xab00
	NBD_SET_BLKSIZE     = 0xab01
	NBD_SET_SIZE        = 0xab02
	NBD_DO_IT           = 0xab03
	NBD_CLEAR_SOCK      = 0xab04
	NBD_SET_FLAGS       = 0xab0a
	NBD_FLAG_SEND_FLUSH = 1 << 2
)

type nbdRequest struct {
	Magic  uint32
	Cmd    uint32
	Handle [8]byte
	From   uint64
	Len    uint32
}

type nbdReply struct {
	Magic  uint32
	Error  uint32
	Handle [8]byte
}

func ioctl(fd uintptr, cmd, arg uintptr) {
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, fd, cmd, arg)
	if errno != 0 {
		log.Fatalf("ioctl 0x%x failed: %v", cmd, errno)
	}
}

func main() {
	conn, err := rados.NewConn()
	if err != nil {
		log.Fatalf("Rados init failed: %v", err)
	}
	if err := conn.ReadDefaultConfigFile(); err != nil {
		log.Fatalf("ReadDefaultConfigFile failed: %v", err)
	}
	if err := conn.Connect(); err != nil {
		log.Fatalf("Connect failed: %v", err)
	}
	defer conn.Shutdown()

	ioctx, err := conn.OpenIOContext(poolName)
	if err != nil {
		log.Fatalf("OpenIOContext failed: %v", err)
	}
	defer ioctx.Destroy()

	image, err := rbd.OpenImage(ioctx, imageName, "")
	if err != nil {
		log.Fatalf("OpenImage failed: %v", err)
	}
	defer image.Close()

	size, err := image.GetSize()
	if err != nil {
		log.Fatalf("GetSize failed: %v", err)
	}

	nbdFile, err := os.OpenFile(nbdDevice, os.O_RDWR, 0)
	if err != nil {
		log.Fatalf("Open NBD device failed: %v", err)
	}
	defer nbdFile.Close()
	nbdFd := nbdFile.Fd()

	sv, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM, 0)
	if err != nil {
		log.Fatalf("Socketpair failed: %v", err)
	}
	defer syscall.Close(sv[0])
	defer syscall.Close(sv[1])

	ioctl(nbdFd, NBD_SET_SIZE, uintptr(size))
	ioctl(nbdFd, NBD_SET_FLAGS, uintptr(NBD_FLAG_SEND_FLUSH))
	ioctl(nbdFd, NBD_SET_SOCK, uintptr(sv[0]))

	go func() {
		sock := os.NewFile(uintptr(sv[1]), "sock")
		defer sock.Close()

		for {
			buf := make([]byte, 28)
			_, err := io.ReadFull(sock, buf)
			if err != nil {
				log.Println("ReadFull error:", err)
				reply := nbdReply{
					Magic: NBD_REPLY_MAGIC,
					Error: 1,
				}
				writeReply(sock, reply)
				continue
			}

			req := nbdRequest{
				Magic: binary.BigEndian.Uint32(buf[0:4]),
				Cmd:   binary.BigEndian.Uint32(buf[4:8]),
				From:  binary.BigEndian.Uint64(buf[16:24]),
				Len:   binary.BigEndian.Uint32(buf[24:28]),
			}
			copy(req.Handle[:], buf[8:16])

			reply := nbdReply{
				Magic:  NBD_REPLY_MAGIC,
				Error:  0,
				Handle: req.Handle,
			}

			if req.Magic != NBD_REQUEST_MAGIC {
				log.Printf("Invalid magic: 0x%x", req.Magic)
				reply.Error = 1
				writeReply(sock, reply)
				continue
			}

			if uint64(req.From)+uint64(req.Len) > size {
				reply.Error = 1
				writeReply(sock, reply)
				continue
			}

			switch req.Cmd {
			case NBD_CMD_READ:
				log.Printf("NBD_CMD_READ offset=%d length=%d", req.From, req.Len)
				data := make([]byte, req.Len)
				n, err := image.ReadAt(data, int64(req.From))
				if err != nil || n != int(req.Len) {
					reply.Error = 1
				}
				writeReply(sock, reply)
				sock.Write(data)

			case NBD_CMD_WRITE:
				log.Printf("NBD_CMD_WRITE offset=%d length=%d", req.From, req.Len)
				data := make([]byte, req.Len)
				_, err := io.ReadFull(sock, data)
				if err != nil {
					reply.Error = 1
				} else {
					go func() {
						_, err = image.WriteAt(data, int64(req.From))
						if err != nil {
							reply.Error = 1
						}
					}()
				}
				writeReply(sock, reply)

			case NBD_CMD_FLUSH:
				log.Println("NBD_CMD_FLUSH")
				if err := image.Flush(); err != nil {
					reply.Error = 1
				}
				writeReply(sock, reply)

			case NBD_CMD_DISC:
				log.Println("NBD_CMD_DISC")
				writeReply(sock, reply)
				return

			case NBD_CMD_TRIM:
				log.Println("NBD_CMD_TRIM")
				_, err := image.Discard(req.From, uint64(req.Len))
				if err != nil {
					reply.Error = 1
				}
				writeReply(sock, reply)

			default:
				log.Printf("Unknown CMD: %d", req.Cmd)
				reply.Error = 1
				writeReply(sock, reply)
			}
		}
	}()

	ioctl(nbdFd, NBD_DO_IT, 0)
	ioctl(nbdFd, NBD_CLEAR_SOCK, 0)
}

func writeReply(w io.Writer, reply nbdReply) {
	buf := make([]byte, 16)
	binary.BigEndian.PutUint32(buf[0:4], reply.Magic)
	binary.BigEndian.PutUint32(buf[4:8], reply.Error)
	copy(buf[8:16], reply.Handle[:])
	w.Write(buf)
}
